from BPR.InteractScore import InteractScore
import numpy as np
import torch
import torch.utils.data as Data
from BPR.FileLoader import FileLoader
from BPR.InteractScore import InteractScore
from tqdm import tqdm
from typing import Dict, List, Tuple


class DataLoader:
    def __init__(self, hyper_params: dict):
        self.hyper_params = hyper_params

        # Load from file
        self.file_loader = FileLoader(hyper_params)
        self.interact_score = InteractScore(hyper_params)

        self.neg_sample_cnt = hyper_params['neg_sample_cnt']
        self.epsilon = hyper_params['epsilon']
        self.show_info = {}

    def generate_user_show_info(self, user_id: int) -> Dict[int, Dict[str, str]]:
        file_loader = FileLoader(self.hyper_params.copy())
        assert user_id in file_loader.intereaction_list

        _, _, show_info = self.interact_score.get_interact_score(
            file_loader.intereaction_list[user_id])
        return show_info

    def generate_dataset(self, save_dataset=True) -> Data.Dataset:
        try:
            dataset, self.show_info = torch.load('BPR/model_dat/dataset.pkl')
            print('Restore from model_dat')
            return dataset
        except Exception as err:
            print(err)

        print('Generating Dataset...')
        positive_item = []
        negative_item = []
        item_weight = []
        user_id = []
        positive_extra = []
        negative_extra = []

        for uid, interact_hist in tqdm(self.file_loader.intereaction_list.items()):
            # Get Positive Negative Pairs
            score_dict, extra_dict, show_info = self.interact_score.get_interact_score(
                interact_hist)
            self.show_info[uid] = show_info

            score_list = sorted(score_dict.items(),
                                key=lambda x: x[1], reverse=True)

            # Phase 1
            book_set = set(score_dict.keys())
            not_book_set = set(range(self.file_loader.book_cnt)) - book_set
            book_len = len(score_list)

            for i in range(book_len-1):
                for j in range(i+1, book_len):
                    if score_list[i][1] - score_list[j][1] > self.epsilon:
                        # Add To Ds
                        positive_item.append(score_list[i][0])
                        negative_item.append(score_list[j][0])
                        item_weight.append(score_list[i][1] - score_list[j][1])
                        user_id.append(uid)
                        positive_extra.append(extra_dict[score_list[i][0]])
                        negative_extra.append(extra_dict[score_list[j][0]])

            # Phase 2
            book_list = []
            for book in book_set:
                if score_dict[book] > self.epsilon:
                    book_list.append(book)

            not_book_list = list(not_book_set)

            phase2_positive = np.random.choice(
                book_list, size=self.neg_sample_cnt)
            phase2_negative = np.random.choice(
                not_book_list, size=self.neg_sample_cnt)
            phase2_weight = list(map(score_dict.get, phase2_positive))

            positive_item.extend(phase2_positive)
            negative_item.extend(phase2_negative)
            item_weight.extend(phase2_weight)
            user_id.extend([uid] * self.neg_sample_cnt)
            positive_extra.extend([extra_dict[i] for i in phase2_positive])
            negative_extra.extend(
                [torch.zeros(self.hyper_params['extra_dim'], dtype=torch.float32)]*self.neg_sample_cnt)

        positive_tensor = torch.LongTensor(positive_item)
        negative_tensor = torch.LongTensor(negative_item)
        weight_tensor = torch.FloatTensor(item_weight)
        user_id_tensor = torch.LongTensor(user_id)
        positive_extra_tensor = torch.stack(positive_extra, dim=0)
        negative_extra_tensor = torch.stack(negative_extra, dim=0)

        dataset = Data.TensorDataset(
            positive_tensor, negative_tensor, weight_tensor, user_id_tensor, positive_extra_tensor, negative_extra_tensor)

        if save_dataset:
            torch.save((dataset, self.show_info), 'BPR/model_dat/dataset.pkl')
        print('Finished!')

        return dataset
